{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f0aa3eec-9b5b-447d-8e38-57b35fd562ec",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# StyleGAN 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "700a2f47-a1fe-44c7-92a2-37a9dc45fc15",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a974642e-ea68-4576-86d3-1a84ca0af73c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install einops ninja gdown click scipy tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d9ae5e0-b8e9-4d6b-a361-4f2b03358af7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Define the static dataset information\n",
    "DATASET = {\n",
    "    \"name\": \"project_name\",\n",
    "    \"path\": \"path/to/dataset\",\n",
    "    \"res\": \"512x512\" # What it will preprocess to\n",
    "}\n",
    "\n",
    "# Function to calculate kimg based on epochs\n",
    "def calculate_kimg(num_images, epochs):\n",
    "    return (num_images * epochs) // 1000\n",
    "\n",
    "# Define the number of images to test\n",
    "NUM_IMAGES = [500, 1000, 2500]\n",
    "\n",
    "# Define the fixed number of epochs\n",
    "EPOCHS = 100\n",
    "\n",
    "training_configs = []\n",
    "for num_images in NUM_IMAGES:\n",
    "    kimg = calculate_kimg(num_images, EPOCHS)\n",
    "    training_configs.append({\n",
    "        \"num_images\": num_images,\n",
    "        \"kimg\": kimg\n",
    "    })\n",
    "\n",
    "# Print the configurations for verification\n",
    "print(\"Training Configurations:\")\n",
    "for config in training_configs:\n",
    "    print(f\"Number of Images: Epochs: {EPOCHS}, {config['num_images']}, kimg: {config['kimg']}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d614de6-d600-403a-a604-171acd7440c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Define the path to the directory\n",
    "dir_path = f\"Results/{DATASET['name']}/model\"\n",
    "\n",
    "os.makedirs(dir_path, exist_ok=True)\n",
    "\n",
    "# Check if the directory is empty\n",
    "if not os.listdir(dir_path):\n",
    "    %cd Results/{DATASET['name']}/model/\n",
    "\n",
    "    !git clone https://github.com/NVlabs/stylegan3\n",
    "\n",
    "# Prepare the path using f-string\n",
    "stylegan3_path = f\"Results/{DATASET['name']}/model/stylegan3\"\n",
    "\n",
    "# Use the formatted path in the %cd command\n",
    "%cd {stylegan3_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9cdf955-5abf-4d01-9b09-33d51e541a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import os\n",
    "import random\n",
    "\n",
    "# Parameters\n",
    "DATASET_PATH = \"path/to/dataset\"\n",
    "OUTPUT_DIR = \"Results/dataset_name/preprocessed\"\n",
    "\n",
    "# Ensure the output directory exists\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "def get_image_paths(directory_path):\n",
    "    image_paths = []\n",
    "    for root, _, files in os.walk(directory_path):\n",
    "        for file in files:\n",
    "            if file.lower().endswith(('.png', '.jpeg', '.jpg')):\n",
    "                image_paths.append(os.path.join(root, file))\n",
    "    return image_paths\n",
    "\n",
    "def sample_and_convert_images(image_paths, num_samples, output_dir):\n",
    "    sampled_paths = random.sample(image_paths, min(num_samples, len(image_paths)))\n",
    "    for image_path in sampled_paths:\n",
    "        filename = os.path.basename(image_path)\n",
    "        output_path = os.path.join(output_dir, filename)\n",
    "        image = Image.open(image_path)\n",
    "        image = image.convert(\"RGB\")\n",
    "        image.save(output_path)\n",
    "        print(f\"Converted {filename} to RGB and saved to {output_path}\")\n",
    "\n",
    "# Get all image paths from the dataset\n",
    "image_paths = get_image_paths(DATASET_PATH)\n",
    "\n",
    "# Specify the number of images to sample\n",
    "num_samples = 2500  # Adjust this number as needed\n",
    "\n",
    "# Sample and convert images\n",
    "sample_and_convert_images(image_paths, num_samples, OUTPUT_DIR)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f3edc9-4b3d-4bd8-8948-3c081bb7f98b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 500 Images "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37409f49-c85e-4c7c-9917-2312f54adbf9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d988c413-e46b-4d08-a2d3-9dab9e9ef6e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = training_configs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1506e12-0109-438b-9733-01bd6f58fbf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "\n",
    "# Directory containing all preprocessed images\n",
    "PREPROCESSED_DIR = f\"Results/{DATASET['name']}/preprocessed\"\n",
    "# Temporary directory to store the first 500 images\n",
    "TEMP_DIR = f\"Results/{DATASET['name']}/{config['num_images']}_preprocessed\"\n",
    "\n",
    "# Ensure the temporary directory exists\n",
    "os.makedirs(TEMP_DIR, exist_ok=True)\n",
    "\n",
    "# Get all image filenames in the preprocessed directory\n",
    "all_images = [f for f in os.listdir(PREPROCESSED_DIR) if f.lower().endswith(('.png', '.jpeg', '.jpg'))]\n",
    "\n",
    "# Sort and select the first X images\n",
    "divided_images = sorted(all_images)[:config['num_images']]\n",
    "\n",
    "# Copy the first 500 images to the temporary directory\n",
    "for image in divided_images:\n",
    "    src_path = os.path.join(PREPROCESSED_DIR, image)\n",
    "    dest_path = os.path.join(TEMP_DIR, image)\n",
    "    shutil.copy(src_path, dest_path)\n",
    "    print(f\"Copied {image} to {TEMP_DIR}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "132c572c-3a41-46dc-bca4-8ab09ebd759c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dest_file = f\"Results/{DATASET['name']}/{config['num_images']}_images/preprocessed-{DATASET['res']}.zip\"\n",
    "\n",
    "!python path/to/dataset_tool.py \\\n",
    "    --source='{TEMP_DIR}' \\\n",
    "    --dest='{dest_file}' \\\n",
    "    --resolution={DATASET['res']}\n",
    "\n",
    "# Delete the temporary directory after processing\n",
    "shutil.rmtree(TEMP_DIR)\n",
    "print(f\"Temporary directory {TEMP_DIR} deleted.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "552abd7c-7856-4504-8eed-308f38aed890",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoints_dir = f\"Results/{DATASET['name']}/{config['num_images']}_images/checkpoints\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "os.makedirs(checkpoints_dir, exist_ok=True)\n",
    "\n",
    "training_data = f\"Results/{DATASET['name']}/{config['num_images']}_images/preprocessed-{DATASET['res']}.zip\"\n",
    "\n",
    "# Start training with adjusted parameters\n",
    "!python path/to/train.py \\\n",
    "  --outdir='{checkpoints_dir}' --cfg=stylegan3-t \\\n",
    "  --data='{training_data}' \\\n",
    "  --snap=10 --metrics=fid50k_full --gpus=3 --batch=96 --batch-gpu=8 --gamma=8 --mirror=1 --kimg={config[\"kimg\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c0da6f6-d9ee-4f05-8692-096337b7cc20",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Generating Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d868dabe-76cd-46fd-88bb-164f6d7869c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_images_dir = f\"Results/{DATASET['name']}/{config['num_images']}_images/generated_images\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "os.makedirs(generated_images_dir, exist_ok=True)\n",
    "\n",
    "gen_network = f\"Results/{DATASET['name']}/{config['num_images']}_images/checkpoints/00001-stylegan3-t-preprocessed-512x512-gpus3-batch96-gamma8/network-snapshot-0000{config['kimg']}.pkl\"\n",
    "\n",
    "# Generate images using the trained model\n",
    "!python path/to/gen_images.py \\\n",
    "  --outdir='{generated_images_dir}' \\\n",
    "  --trunc=0.7 --seeds=0-24 \\\n",
    "  --network='{gen_network}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf6b6edb-6e5f-42de-b24b-4de1b75d7db5",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 1000 Images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fdf026c-5207-43d8-b27c-b639100a381f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59851cfe-4630-4193-8535-cb1e5e2e26d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = training_configs[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c68a309-614e-4eeb-9371-f46dfa468c92",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "\n",
    "# Directory containing all preprocessed images\n",
    "PREPROCESSED_DIR = f\"Results/{DATASET['name']}/preprocessed\"\n",
    "# Temporary directory to store the first 500 images\n",
    "TEMP_DIR = f\"Results/{DATASET['name']}/{config['num_images']}_preprocessed\"\n",
    "\n",
    "# Ensure the temporary directory exists\n",
    "os.makedirs(TEMP_DIR, exist_ok=True)\n",
    "\n",
    "# Get all image filenames in the preprocessed directory\n",
    "all_images = [f for f in os.listdir(PREPROCESSED_DIR) if f.lower().endswith(('.png', '.jpeg', '.jpg'))]\n",
    "\n",
    "# Sort and select the first X images\n",
    "divided_images = sorted(all_images)[:config['num_images']]\n",
    "\n",
    "# Copy the first X images to the temporary directory\n",
    "for image in divided_images:\n",
    "    src_path = os.path.join(PREPROCESSED_DIR, image)\n",
    "    dest_path = os.path.join(TEMP_DIR, image)\n",
    "    shutil.copy(src_path, dest_path)\n",
    "    print(f\"Copied {image} to {TEMP_DIR}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "851320da-0504-40f3-81ad-881e45a019a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "dest_file = f\"Results/{DATASET['name']}/{config['num_images']}_images/preprocessed-{DATASET['res']}.zip\"\n",
    "\n",
    "!python path/to/dataset_tool.py \\\n",
    "    --source='{TEMP_DIR}' \\\n",
    "    --dest='{dest_file}' \\\n",
    "    --resolution={DATASET['res']}\n",
    "\n",
    "# Delete the temporary directory after processing\n",
    "shutil.rmtree(TEMP_DIR)\n",
    "print(f\"Temporary directory {TEMP_DIR} deleted.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d53eb5e-3b31-4851-af25-770f064b6a9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoints_dir = f\"Results/{DATASET['name']}/{config['num_images']}_images/checkpoints\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "os.makedirs(checkpoints_dir, exist_ok=True)\n",
    "\n",
    "training_data = f\"Results/{DATASET['name']}/{config['num_images']}_images/preprocessed-{DATASET['res']}.zip\"\n",
    "\n",
    "# Start training with adjusted parameters\n",
    "!python path/to/train.py \\\n",
    "  --outdir='{checkpoints_dir}' --cfg=stylegan3-t \\\n",
    "  --data='{training_data}' \\\n",
    "  --snap=10 --metrics=fid50k_full --gpus=3 --batch=24 --batch-gpu=8 --gamma=8 --mirror=1 --kimg={config[\"kimg\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6537cd0-0c08-4847-92cb-6f60ff98961a",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Generating Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b865412-d62d-49ef-a003-9a25441e24f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_images_dir = f\"Results/{DATASET['name']}/{config['num_images']}_images/generated_images\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "os.makedirs(generated_images_dir, exist_ok=True)\n",
    "\n",
    "gen_network = f\"Results/{DATASET['name']}/{config['num_images']}_images/checkpoints/00002-stylegan3-t-preprocessed-512x512-gpus3-batch24-gamma8/network-snapshot-000{config['kimg']}.pkl\"\n",
    "\n",
    "# Generate images using the trained model\n",
    "!python path/to/gen_images.py \\\n",
    "  --outdir='{generated_images_dir}' \\\n",
    "  --trunc=1 --seeds=0-24 \\\n",
    "  --network='{gen_network}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6ccc206-0b37-4bb1-be02-d426e40cc1a1",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## 2500 Images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf4c9e5-95de-4cb8-9fde-be4b5aa2b763",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a335d409-332f-4007-a559-717f3adba2a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = training_configs[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92750e89-2b3f-4383-92fe-695060fef349",
   "metadata": {},
   "outputs": [],
   "source": [
    "dest_file = f\"Results/{DATASET['name']}/{config['num_images']}_images/preprocessed-{DATASET['res']}.zip\"\n",
    "FULL_PREPROCESSED = f\"Results/{DATASET['name']}/preprocessed\"\n",
    "\n",
    "!python path/to/dataset_tool.py \\\n",
    "    --source='{FULL_PREPROCESSED}' \\\n",
    "    --dest='{dest_file}' \\\n",
    "    --resolution={DATASET['res']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e99a97b-595a-432f-9260-9209b008e953",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoints_dir = f\"Results/{DATASET['name']}/{config['num_images']}_images/checkpoints\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "os.makedirs(checkpoints_dir, exist_ok=True)\n",
    "\n",
    "training_data = f\"Results/{DATASET['name']}/{config['num_images']}_images/preprocessed-{DATASET['res']}.zip\"\n",
    "\n",
    "# Start training with adjusted parameters\n",
    "!python path/to/train.py \\\n",
    "  --outdir='{checkpoints_dir}' --cfg=stylegan3-t \\\n",
    "  --data='{training_data}' \\\n",
    "  --snap=10 --metrics=fid50k_full --gpus=3 --batch=24 --batch-gpu=8 --gamma=8 --mirror=1 --kimg={config[\"kimg\"]} \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ffbfec4-4bf7-4001-a73a-d3645ca8ae42",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Generating Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52bf77a8-790e-4145-a3b1-52eead579687",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_images_dir = f\"Results/{DATASET['name']}/{config['num_images']}_images/generated_images\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "os.makedirs(generated_images_dir, exist_ok=True)\n",
    "\n",
    "# MAKE SURE TO CHANGE MANUALLY TO CORRECT PATH\n",
    "gen_network = f\"Results/{DATASET['name']}/{config['num_images']}_images/checkpoints/00004-stylegan3-t-preprocessed-512x512-gpus3-batch24-gamma8/network-snapshot-000{config['kimg']}.pkl\"\n",
    "\n",
    "# Generate images using the trained model\n",
    "!python path/to/gen_images.py \\\n",
    "  --outdir='{generated_images_dir}' \\\n",
    "  --trunc=0.5 --seeds=0-24 \\\n",
    "  --network='{gen_network}'"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
